#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
添加身份证号唯一性约束的数据库迁移脚本
"""

import sqlite3
import os
from datetime import datetime

def add_id_card_unique_constraint():
    """为身份证号字段添加唯一性约束"""
    
    # 数据库文件路径
    db_path = os.path.join(os.path.dirname(__file__), '..', 'instance', 'dispatch_system.db')
    
    if not os.path.exists(db_path):
        print(f"数据库文件不存在: {db_path}")
        return False
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        print("开始添加身份证号唯一性约束...")
        
        # 检查是否已经有唯一性约束
        cursor.execute("PRAGMA table_info(users)")
        columns = cursor.fetchall()
        
        # 检查是否存在重复的身份证号
        cursor.execute("""
            SELECT id_card, COUNT(*) as count 
            FROM users 
            WHERE id_card IS NOT NULL AND id_card != '' 
            GROUP BY id_card 
            HAVING COUNT(*) > 1
        """)
        duplicates = cursor.fetchall()
        
        if duplicates:
            print("发现重复的身份证号:")
            for id_card, count in duplicates:
                print(f"  身份证号: {id_card}, 重复次数: {count}")
            
            # 为重复的身份证号添加后缀
            for id_card, count in duplicates:
                cursor.execute("""
                    SELECT id, username FROM users 
                    WHERE id_card = ? 
                    ORDER BY id
                """, (id_card,))
                users_with_duplicate = cursor.fetchall()
                
                # 保留第一个用户的身份证号不变，其他用户添加后缀
                for i, (user_id, username) in enumerate(users_with_duplicate[1:], 1):
                    new_id_card = f"{id_card}_dup_{i}"
                    cursor.execute("""
                        UPDATE users 
                        SET id_card = ? 
                        WHERE id = ?
                    """, (new_id_card, user_id))
                    print(f"  用户 {username} (ID: {user_id}) 的身份证号已更新为: {new_id_card}")
        
        # 由于SQLite的限制，我们需要重新创建表来添加唯一性约束
        # 首先检查并删除可能存在的备份表
        cursor.execute("DROP TABLE IF EXISTS users_backup")
        
        # 备份现有数据
        cursor.execute("""
            CREATE TABLE users_backup AS 
            SELECT * FROM users
        """)
        
        # 删除原表
        cursor.execute("DROP TABLE users")
        
        # 首先更新空值字段
        cursor.execute("""
            UPDATE users_backup 
            SET 
                name = COALESCE(name, '未知'),
                employee_id = COALESCE(employee_id, 'EMP' || id),
                phone = COALESCE(phone, ''),
                position = COALESCE(position, '未知'),
                department = COALESCE(department, '未知'),
                area = COALESCE(area, '未知'),
                id_card = COALESCE(id_card, '000000000000000000'),
                ethnicity = COALESCE(ethnicity, '汉族'),
                gender = COALESCE(gender, 'male'),
                hire_date = COALESCE(hire_date, date('now')),
                is_driver = COALESCE(is_driver, 0)
            WHERE 
                name IS NULL OR name = '' OR
                employee_id IS NULL OR employee_id = '' OR
                phone IS NULL OR
                position IS NULL OR position = '' OR
                department IS NULL OR department = '' OR
                area IS NULL OR area = '' OR
                id_card IS NULL OR id_card = '' OR
                ethnicity IS NULL OR ethnicity = '' OR
                gender IS NULL OR gender = '' OR
                hire_date IS NULL OR
                is_driver IS NULL
        """)
        
        # 重新创建表，包含唯一性约束
        cursor.execute("""
            CREATE TABLE users (
                id INTEGER PRIMARY KEY,
                username VARCHAR(80) UNIQUE NOT NULL,
                email VARCHAR(120) UNIQUE NOT NULL,
                password_hash VARCHAR(255) NOT NULL,
                role VARCHAR(20) NOT NULL DEFAULT 'user',
                is_active BOOLEAN DEFAULT 1,
                name VARCHAR(100) NOT NULL,
                employee_id VARCHAR(50) UNIQUE NOT NULL,
                phone VARCHAR(20) NOT NULL,
                position VARCHAR(100) NOT NULL,
                department VARCHAR(100) NOT NULL,
                skills TEXT,
                area VARCHAR(100) NOT NULL,
                id_card VARCHAR(18) UNIQUE NOT NULL,
                ethnicity VARCHAR(50) NOT NULL,
                gender VARCHAR(10) NOT NULL,
                hire_date DATE NOT NULL,
                is_driver BOOLEAN NOT NULL DEFAULT 0,
                password_field VARCHAR(255),
                status VARCHAR(20) DEFAULT 'active',
                resignation_date DATE,
                electrician_cert_expiry DATE,
                height_cert_expiry DATE,
                batch_number VARCHAR(50),
                remarks TEXT,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                group_id INTEGER,
                FOREIGN KEY (group_id) REFERENCES groups (id)
            )
        """)
        
        # 恢复数据
        cursor.execute("""
            INSERT INTO users 
            SELECT * FROM users_backup
        """)
        
        # 删除备份表
        cursor.execute("DROP TABLE users_backup")
        
        conn.commit()
        print("身份证号唯一性约束添加成功！")
        
        return True
        
    except Exception as e:
        print(f"添加唯一性约束时出错: {e}")
        conn.rollback()
        return False
    finally:
        conn.close()

if __name__ == '__main__':
    print(f"开始执行数据库迁移 - {datetime.now()}")
    success = add_id_card_unique_constraint()
    if success:
        print("迁移完成！")
    else:
        print("迁移失败！")